# 1、读取图片的数据
# 特征工程
import os
import cv2
import numpy as np
import pandas as pd

image_base_url = "C:\\Users\\shujia\\Desktop\\train\\"

images_data = pd.DataFrame(columns=["image_name", "image_data"])

# 循环读取目录下所有的图片
image_names = os.listdir(image_base_url)
for image_name in image_names:
    # 拼接图片的路径
    image_url = f"{image_base_url}{image_name}"
    # 读取图片
    image_data = cv2.imread(image_url, 0)
    # 将图片矩阵转换向量
    image_data = image_data.reshape(784)
    # 归一化
    image_data = [1 if i > 126 else 0 for i in image_data]
    # 将数据保存到df中
    images_data.loc[len(images_data)] = [image_name, image_data]

print("图片读取完成")

# 读取图片数字标记文件
image_label = pd.read_csv("../data/image_res.txt", sep=" ", names=["image_name", "y"])

# 关联获取图片的数字
image_label_data = pd.merge(images_data, image_label, on="image_name")

# 取出x和y,切分数据集
y = image_label_data["y"]
image_data = image_label_data["image_data"]
# 将数据转换成ndArray
x = np.array([np.array(line) for line in image_data])

# 切分数据集
from sklearn.model_selection import train_test_split

train_x, text_x, train_y, test_y = train_test_split(x, y, test_size=0.2)

# 选择算法训练模型
from sklearn.linear_model import LogisticRegression

lr = LogisticRegression()
# 训练模型
model = lr.fit(train_x, train_y)

print(f"模型准确率:{model.score(text_x, test_y)}")

# 保存模型

import pickle

# 模型的保存
with open('图片分类模型.pickle', 'wb') as f:
    pickle.dump(lr, f)  # 将训练好的模型clf存储在变量f中，且保存到本地
